import os
import numpy as np
from skimage.metrics import structural_similarity as ssim
from skimage.metrics import peak_signal_noise_ratio as psnr
from scipy.stats import pearsonr
import cv2
from datetime import datetime


def evaluate_models_with_save():
    """
    评估生成图片质量，并将结果保存到txt文件
    """
    # 1. 配置路径
    real_images_dir = r'I:\lxj\stainedHE\pytorch_CycleGAN_and_pix2pix_master\virtual_staining_eval\data\real_images_diffusion'
    fake_images_dir = r'I:\lxj\stainedHE\pytorch_CycleGAN_and_pix2pix_master\virtual_staining_eval\data\fake_images_diffusion'
    results_dir = r'I:\lxj\stainedHE\pytorch_CycleGAN_and_pix2pix_master\virtual_staining_eval\results'

    # 创建结果文件夹
    os.makedirs(results_dir, exist_ok=True)

    # 生成结果文件名（包含时间戳）
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_filename = f"evaluation_results_{timestamp}.txt"
    result_filepath = os.path.join(results_dir, result_filename)

    # 2. 获取图像文件列表
    real_images_list = sorted(
        [f for f in os.listdir(real_images_dir) if f.endswith(('.png', '.jpg', '.jpeg', '.tif', '.tiff', '.bmp'))])
    fake_images_list = sorted(
        [f for f in os.listdir(fake_images_dir) if f.endswith(('.png', '.jpg', '.jpeg', '.tif', '.tiff', '.bmp'))])

    # 确保文件列表对应
    assert len(real_images_list) == len(fake_images_list), "真实图像和生成图像数量必须相同"

    # 3. 初始化存储结果的列表
    ssim_scores = []
    psnr_scores = []
    pcc_scores = []

    # 4. 打开文件准备写入
    with open(result_filepath, 'w', encoding='utf-8') as f:
        # 写入文件头信息
        f.write("=" * 70 + "\n")
        f.write("H&E染色图像质量评估报告\n")
        f.write("=" * 70 + "\n")
        f.write(f"评估时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"真实图像目录: {real_images_dir}\n")
        f.write(f"生成图像目录: {fake_images_dir}\n")
        f.write(f"评估图像对数: {len(real_images_list)}\n")
        f.write("=" * 70 + "\n\n")

        # 5. 遍历每一对图像
        f.write("逐图像评估结果:\n")
        f.write("-" * 70 + "\n")

        for real_img_name, fake_img_name in zip(real_images_list, fake_images_list):
            real_img_path = os.path.join(real_images_dir, real_img_name)
            fake_img_path = os.path.join(fake_images_dir, fake_img_name)

            # 读取图像
            real_img = cv2.imread(real_img_path, cv2.IMREAD_COLOR)
            fake_img = cv2.imread(fake_img_path, cv2.IMREAD_COLOR)

            # 确保图像读取成功
            if real_img is None or fake_img is None:
                error_msg = f"警告：无法读取 {real_img_path} 或 {fake_img_path}，跳过。\n"
                print(error_msg)
                f.write(error_msg)
                continue

            # 确保图像数据类型为浮点数，范围[0, 1]
            real_img = real_img.astype(np.float64) / 255.0
            fake_img = fake_img.astype(np.float64) / 255.0

            # 计算SSIM
            ssim_value = ssim(real_img, fake_img, win_size=7, channel_axis=-1, data_range=1.0)
            ssim_scores.append(ssim_value)

            # 计算PSNR
            psnr_value = psnr(real_img, fake_img, data_range=1.0)
            psnr_scores.append(psnr_value)

            # 计算PCC
            real_flat = real_img.flatten()
            fake_flat = fake_img.flatten()
            pcc_value, _ = pearsonr(real_flat, fake_flat)
            pcc_scores.append(pcc_value)

            # 打印并写入每对图像的结果
            result_line = f"{real_img_name} / {fake_img_name}: SSIM={ssim_value:.4f}, PSNR={psnr_value:.4f} dB, PCC={pcc_value:.4f}\n"
            print(result_line.strip())  # 控制台打印
            f.write(result_line)  # 文件写入

        # 6. 计算总体统计信息
        f.write("\n" + "=" * 70 + "\n")
        f.write("评估结果总结:\n")
        f.write("=" * 70 + "\n")

        if ssim_scores:  # 确保有有效数据
            # 基本统计
            summary_stats = [
                f"评估图像对数: {len(ssim_scores)}",
                f"SSIM 平均值 ± 标准差: {np.mean(ssim_scores):.4f} ± {np.std(ssim_scores):.4f}",
                f"SSIM 范围: [{np.min(ssim_scores):.4f}, {np.max(ssim_scores):.4f}]",
                f"PSNR 平均值 ± 标准差: {np.mean(psnr_scores):.4f} ± {np.std(psnr_scores):.4f} dB",
                f"PSNR 范围: [{np.min(psnr_scores):.4f}, {np.max(psnr_scores):.4f}] dB",
                f"PCC 平均值 ± 标准差: {np.mean(pcc_scores):.4f} ± {np.std(pcc_scores):.4f}",
                f"PCC 范围: [{np.min(pcc_scores):.4f}, {np.max(pcc_scores):.4f}]"
            ]

            for stat in summary_stats:
                print(stat)
                f.write(stat + "\n")

            # 详细数据（可选）
            f.write("\n详细数据:\n")
            f.write("-" * 70 + "\n")
            f.write("SSIM values: " + ", ".join([f"{x:.4f}" for x in ssim_scores]) + "\n")
            f.write("PSNR values: " + ", ".join([f"{x:.4f}" for x in psnr_scores]) + "\n")
            f.write("PCC values: " + ", ".join([f"{x:.4f}" for x in pcc_scores]) + "\n")

        else:
            error_msg = "错误：没有有效的图像对可用于评估。"
            print(error_msg)
            f.write(error_msg + "\n")

        # 7. 文件尾信息
        f.write("\n" + "=" * 70 + "\n")
        f.write("评估完成\n")
        f.write(f"结果文件保存至: {result_filepath}\n")
        f.write("=" * 70 + "\n")

    # 8. 在控制台也显示最终信息
    print("\n" + "=" * 70)
    print("评估完成！")
    print(f"详细结果已保存至: {result_filepath}")
    print("=" * 70)


if __name__ == "__main__":
    evaluate_models_with_save()